# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/11 20:22 
@Author  : ZhangShenao 
@File    : 2.使用FewShotPromptTemplate少样本提示词.py 
@Desc    : 使用FewShotPromptTemplate,实现少样本提示词

使用FewShotPromptTemplate模板,可以构建少样本提示词

使用SemanticSimilarityExampleSelector示例选择器,可以通过语义上的余弦相似度,选择出最合适的示例作为样本,发送给LLM,以节约Token的使用量
"""

import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate, ChatPromptTemplate

dotenv.load_dotenv()

# 预设示例
# 1. 创建一些示例
samples = [
    {
        "flower_type": "玫瑰",
        "occasion": "爱情",
        "ad_copy": "玫瑰，浪漫的象征，是你向心爱的人表达爱意的最佳选择。"
    },
    {
        "flower_type": "康乃馨",
        "occasion": "母亲节",
        "ad_copy": "康乃馨代表着母爱的纯洁与伟大，是母亲节赠送给母亲的完美礼物。"
    },
    {
        "flower_type": "百合",
        "occasion": "庆祝",
        "ad_copy": "百合象征着纯洁与高雅，是你庆祝特殊时刻的理想选择。"
    },
    {
        "flower_type": "向日葵",
        "occasion": "鼓励",
        "ad_copy": "向日葵象征着坚韧和乐观，是你鼓励亲朋好友的最好方式。"
    }
]

# 创建示例提示模板
sample_prompt = PromptTemplate(
    template='鲜花类型：{flower_type}\n场合：{occasion}\n广告文案：{ad_copy}',
    input_variables=['flower_type', 'occasion', 'ad_copy']
)

# 创建示例选择器,根据向量相似度,选择合适的示例,而不是把全部示例都发给LLM
# 使用示例选择器,可以通过语义上的余弦相似度,选择出最合适的示例作为样本,发送给LLM,以节约Token的使用量
example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples=samples,  # 指定示例样本
    embeddings=ZhipuAIEmbeddings(),  # 指定嵌入模型
    vectorstore_cls=Chroma,  # 指定向量数据库
    k=1,  # 指定最终选择的示例数量
)

# 创建FewShotPromptTemplate少样本提示模板
few_shot_prompt = FewShotPromptTemplate(
    example_selector=example_selector,  # 指定示例选择器
    example_prompt=sample_prompt,  # 指定示例提示词
    suffix='请根据提供的示例，生成一个广告文案：\n鲜花类型：{flower_type}\n场合：{occasion}\n',  # 指定示例后的提示词
    input_variables=['flower_type', 'occasion'],  # 指定输入变量
)

# 打印Few Shot Prompt
# 向量数据库会根据余弦相似度,选择'百合'这个示例
# print(few_shot_prompt.format(flower_type='白百何', occasion='生日'))

# 构建聊天Prompt
prompt = ChatPromptTemplate.from_messages([
    ('system', '你是一个强大的聊天机器人，请根据用户的提问生成回答'),
    ('human', '{query}')
])

# 构造Chain
llm = ChatZhipuAI(model='glm-4-air')
parser = StrOutputParser()
chain = prompt | llm | parser

# 执行Chain,获取结果
output = chain.invoke({'query': few_shot_prompt.format(flower_type='满天星', occasion='生日')})
print(output)
